/*
 * hmm.cc
 *
 *  Created on: Jul 4, 2011
 *      Author: tqlong
 */

#include "hmm.h"

HiddenMarkovModel::HiddenMarkovModel(int n_state, int n_observation) :
  n_state_(n_state), n_observation_(n_observation)
{
  initialize();
}

void
HiddenMarkovModel::initialize()
{
  setStart(arma::randu<vec_type>(n_state_));
  setTransition(arma::randu<mat_type>(n_state_, n_state_));
  setEmission(arma::randu<mat_type>(n_state_, n_observation_));
}

/** Divide each row of the matrix a by the sum of its elements */
void
HiddenMarkovModel::toProbabilityMatrix(mat_type &a)
{
  for (unsigned int i = 0; i < a.n_rows; i++)
    {
      double s = 0;
      for (unsigned int j = 0; j < a.n_cols; j++)
        s += a.at(i, j);
      if (s > 10e-10) a.row(i) /= s;
      else a.row(i).zeros();
    }
}

void
HiddenMarkovModel::dump()
{
  pi_.print("pi=");
  a_.print("a=");
  b_.print("b=");

  logpi_.print("logpi=");
  loga_.print("loga=");
  logb_.print("logb=");
}

void
HiddenMarkovModel::genSeq(int length, seq_type &sseq, seq_type &oseq)
{
  sseq = oseq = seq_type(length);
  int start = randomWheel(pi_);
  for (int t = 0; t < length; t++)
    {
      if (t == 0) sseq[t] = start;
      else sseq[t] = randomWheel(a_.row(sseq[t - 1]));
      oseq[t] = randomWheel(b_.row(sseq[t]));
    }
}

void
HiddenMarkovModel::print_seq(const seq_type &sseq, const seq_type &oseq)
{
  assert(sseq.size() == oseq.size());
  for (unsigned int t = 0; t < sseq.size(); t++)
    std::cout << "(" << sseq[t] << "," << oseq[t] << ") ";
  std::cout << std::endl;
}

double
HiddenMarkovModel::logmul(double a, double b) const
{
  return a + b;
}

double
HiddenMarkovModel::logdiv(double a, double b) const
{
  return a - b;
}

double
HiddenMarkovModel::logadd(double a, double b) const
{
  if (a > b) return a + log1p(exp(b - a));
  else return b + log1p(exp(a - b));
}

double
HiddenMarkovModel::probability(const seq_type &oseq)
{
  seq_ = oseq;
  forward_pass();
  //backward_pass(); // comment if no sanity check needed (see below)

  int T = seq_.size();
  double p = LOGZERO;
  for (int i = 0; i < n_state_; i++)
    p = logadd(p, alpha_.at(T - 1, i));

  // sanity check if backward_pass() is used above
  //  std::cout << "p=" << p << std::endl;
  //  for (int t = 0; t < T; t++)
  //    {
  //      double s = LOGZERO;
  //      for (int i = 0; i < n_state_; i++)
  //        {
  //          double v = logmul(alpha_.at(t, i), beta_.at(t, i));
  //          s = logadd(s, v);
  //        }
  //      std::cout << "s=" << s << std::endl;
  //    }

  return p;
}

void
HiddenMarkovModel::map(const seq_type &oseq, seq_type &sseq)
{
  seq_ = oseq;
  viterbi_pass();

  int T = seq_.size();
  double s = LOGZERO;
  int state = -1;
  for (int i = 0; i < n_state_; i++)
    if (s < alpha_.at(T - 1, i))
      {
        s = alpha_.at(T - 1, i);
        state = i;
      }
  assert(state >= 0);

  // back-tracking using trace
  sseq = oseq;
  sseq[T - 1] = state;
  for (int t = T - 2; t >= 0; t--)
    sseq[t] = trace_.at(t + 1, sseq[t + 1]);
}

void
HiddenMarkovModel::train(const vseq_type& data)
{
  accpi_ = pi_;
  acca_ = a_;
  accb_ = b_;

  for (int iter = 0; iter < 100; iter++)
    {
      initialize_accumulate();
      double likelihood = log(1);
      for (unsigned int i_seq = 0; i_seq < data.size(); i_seq++)
        {
          //seq_ = data[i_seq];
          double p = probability(data[i_seq]);
          backward_pass();
          accumulate();
          likelihood = logmul(likelihood, p);
        }
      reestimate();
      std::cout << "iter=" << iter << " logL=" << likelihood << std::endl;
    }
}

void
HiddenMarkovModel::viterbi_pass()
{
  int T = seq_.size();
  alpha_ = arma::zeros<mat_type>(T, n_state_);
  trace_ = arma::zeros<trace_type>(T, n_state_);

  // initialize t = 0
  for (int i = 0; i < n_state_; i++)
    alpha_.at(0, i) = logmul(logpi_[i], logb_.at(i, seq_[0]));

  for (int t = 1; t < T; t++)
    {
      for (int i = 0; i < n_state_; i++)
        {
          double s = LOGZERO;
          for (int j = 0; j < n_state_; j++)
            {
              double v = logmul(alpha_.at(t - 1, j), loga_.at(j, i));
              if (s < v)
                {
                  s = v;
                  trace_.at(t, i) = j;
                }
            }
          alpha_.at(t, i) = logmul(s, logb_.at(i, seq_[t]));
        }
    }
  //alpha_.print("alpha=");
}

void
HiddenMarkovModel::forward_pass()
{
  int length = seq_.size();
  alpha_ = arma::zeros<mat_type>(length, n_state_);

  // initialize t = 0
  for (int i = 0; i < n_state_; i++)
    alpha_.at(0, i) = logmul(logpi_[i], logb_.at(i, seq_[0]));

  for (int t = 1; t < length; t++)
    {
      for (int i = 0; i < n_state_; i++)
        {
          double s = LOGZERO;
          for (int j = 0; j < n_state_; j++)
            {
              s = logadd(s, logmul(alpha_.at(t - 1, j), loga_.at(j, i)));
            }
          alpha_.at(t, i) = logmul(s, logb_.at(i, seq_[t]));
        }
    }
  //alpha_.print("alpha=");
}

void
HiddenMarkovModel::backward_pass()
{
  int T = seq_.size();
  beta_ = arma::zeros<mat_type>(T, n_state_);

  // initialize t = T-1
  for (int i = 0; i < n_state_; i++)
    beta_.at(T - 1, i) = 0;

  for (int t = T - 2; t >= 0; t--)
    {
      for (int i = 0; i < n_state_; i++)
        {
          double s = LOGZERO;
          for (int j = 0; j < n_state_; j++)
            {
              double v = logmul(beta_.at(t + 1, j), loga_.at(i, j));
              v = logmul(v, logb_.at(j, seq_[t + 1]));
              s = logadd(s, v);
            }
          beta_.at(t, i) = s;
        }
    }
  //beta_.print("beta=");
}

void
HiddenMarkovModel::setTransition(const mat_type &a)
{
  assert(a.n_rows == a.n_cols && (int) a.n_rows == n_state_);
  a_ = a;
  toProbabilityMatrix(a_);
  loga_ = arma::log(a_);
}

void
HiddenMarkovModel::setEmission(const mat_type &b)
{
  assert((int)b.n_rows == n_state_ && (int)b.n_cols == n_observation_);
  b_ = b;
  toProbabilityMatrix(b_);
  logb_ = arma::log(b_);
}

void
HiddenMarkovModel::setStart(const vec_type &pi)
{
  assert((int)pi.n_elem == n_state_);
  pi_ = pi / arma::sum(pi);
  logpi_ = arma::log(pi_);
}

const HiddenMarkovModel::mat_type&
HiddenMarkovModel::transition() const
{
  return a_;
}

const HiddenMarkovModel::mat_type&
HiddenMarkovModel::emission() const
{
  return b_;
}

const HiddenMarkovModel::vec_type&
HiddenMarkovModel::start() const
{
  return pi_;
}

void
HiddenMarkovModel::initialize_accumulate()
{
  accpi_.fill(LOGZERO);
  acca_.fill(LOGZERO);
  accb_.fill(LOGZERO);
}

void
HiddenMarkovModel::accumulate()
{
  // accumulate for pi
  for (int i = 0; i < n_state_; i++)
    {
      double v = logmul(alpha_.at(0, i), beta_.at(0, i));
      accpi_[i] = logadd(accpi_[i], v);
    }

  int T = seq_.size();
  // accumulate for a
  for (int i = 0; i < n_state_; i++)
    for (int j = 0; j < n_state_; j++)
      {
        for (int t = 0; t < T - 1; t++)
          {
            double v = logmul(alpha_.at(t, i), loga_.at(i, j));
            double w = logmul(beta_.at(t + 1, j), logb_.at(j, seq_[t + 1]));
            v = logmul(v, w);
            acca_.at(i, j) = logadd(acca_.at(i, j), v);
          }
      }

  // accumulate for b
  for (int i = 0; i < n_state_; i++)
    for (int t = 0; t < T; t++)
      {
        int k = seq_[t];
        double v = logmul(alpha_.at(t, i), beta_.at(t, i));
        accb_.at(i, k) = logadd(accb_.at(i, k), v);
      }
}

void
HiddenMarkovModel::reestimate()
{
  // pi
  double s = LOGZERO;
  for (int i = 0; i < n_state_; i++)
    s = logadd(s, accpi_[i]);
  for (int i = 0; i < n_state_; i++)
    logpi_[i] = logdiv(accpi_[i], s);
  pi_ = arma::exp(logpi_);

  // a
  for (int i = 0; i < n_state_; i++)
    {
      double s = LOGZERO;
      for (int j = 0; j < n_state_; j++)
        s = logadd(s, acca_.at(i, j));
      for (int j = 0; j < n_state_; j++)
        loga_.at(i, j) = logdiv(acca_.at(i, j), s);
    }
  a_ = arma::exp(loga_);

  // b
  for (int i = 0; i < n_state_; i++)
    {
      double s = LOGZERO;
      for (int k = 0; k < n_observation_; k++)
        s = logadd(s, accb_.at(i, k));
      for (int k = 0; k < n_observation_; k++)
        logb_.at(i, k) = logdiv(accb_.at(i, k), s);
    }
  b_ = arma::exp(logb_);
}

// end of file
